#include <cassert>
#include <chrono>
#include <iostream>
#include <pthread.h>
#include <random>
#include <vector>

std::random_device rd;

class Matrix {
private:
  int rows_;
  int cols_;
  std::vector<double> data_;
  std::mt19937 gen;

public:
  Matrix(int rows, int cols) : rows_(rows), cols_(cols) {
    data_.resize(rows * cols);
    this->gen = std::mt19937(rd());
  }

  void random() {
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < cols_; j++) {
        (*this)(i, j) = std::uniform_real_distribution<double>(0, 1)(gen);
      }
    }
  }

  int rows() const { return rows_; }
  int cols() const { return cols_; }

  double &operator()(int i, int j) { return data_[i * cols_ + j]; }
  double operator()(int i, int j) const { return data_[i * cols_ + j]; }
  bool operator==(const Matrix &B) const {
    if (rows_ != B.rows_ || cols_ != B.cols_) {
      return false;
    }
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < cols_; j++) {
        if ((*this)(i, j) != B(i, j)) {
          return false;
        }
      }
    }
    return true;
  }

  Matrix operator*(const Matrix &B) {
    Matrix C(rows_, B.cols_);
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < B.cols_; j++) {
        C(i, j) = 0;
        for (int k = 0; k < cols_; k++) {
          C(i, j) += (*this)(i, k) * B(k, j);
        }
      }
    }
    return C;
  }

  struct ThreadData {
    const Matrix *A;
    const Matrix *B;
    Matrix *C;
    int start_row;
    int end_row;
  };

  static void *multiplyThread(void *arg) {
    ThreadData *data = (ThreadData *)arg;
    for (int i = data->start_row; i < data->end_row; i++) {
      for (int j = 0; j < data->B->cols(); j++) {
        (*data->C)(i, j) = 0;
        for (int k = 0; k < data->A->cols(); k++) {
          (*data->C)(i, j) += (*data->A)(i, k) * (*data->B)(k, j);
        }
      }
    }
    return nullptr;
  }

  Matrix pthreadMul(const Matrix &B, int num_threads = 4) {
    Matrix C(rows_, B.cols_);
    std::vector<pthread_t> threads(num_threads);
    std::vector<ThreadData> thread_data(num_threads);

    int rows_per_thread = rows_ / num_threads;

    for (int i = 0; i < num_threads; i++) {
      thread_data[i].A = this;
      thread_data[i].B = &B;
      thread_data[i].C = &C;
      thread_data[i].start_row = i * rows_per_thread;
      thread_data[i].end_row =
          (i == num_threads - 1) ? rows_ : (i + 1) * rows_per_thread;

      pthread_create(&threads[i], nullptr, multiplyThread, &thread_data[i]);
    }

    for (int i = 0; i < num_threads; i++) {
      pthread_join(threads[i], nullptr);
    }

    return C;
  }

  void print() {
    for (int i = 0; i < rows_; i++) {
      for (int j = 0; j < cols_; j++) {
        std::cout << (*this)(i, j) << " ";
      }
      std::cout << std::endl;
    }
  }
};

int main(int argc, char **argv) {
  // args: N, Number of threads

  int N = std::stoi(argv[1]);
  int num_threads = std::stoi(argv[2]);

  Matrix A(N, N);
  A.random();
  Matrix B(N, N);
  B.random();

  std::chrono::steady_clock::time_point begin =
      std::chrono::steady_clock::now();
  Matrix C = A * B;
  std::chrono::steady_clock::time_point end = std::chrono::steady_clock::now();
  auto duration =
      std::chrono::duration_cast<std::chrono::microseconds>(end - begin)
          .count();

  begin = std::chrono::steady_clock::now();
  Matrix D = A.pthreadMul(B, num_threads);
  end = std::chrono::steady_clock::now();
  auto duration_pthread =
      std::chrono::duration_cast<std::chrono::microseconds>(end - begin)
          .count();

  std::cout << "Normal: " << duration << std::endl;
  std::cout << "Pthread: " << duration_pthread << std::endl;

  assert(C == D);
  return 0;
}